import argparse
import json
import logging
import os
import sys
import torch
import torchaudio

from denoiser.audio import Audioset, find_audio_files
from denoiser.pretrained import dns48, dns64, master64, valentini_nc
from denoiser.demucs import DemucsStreamer
from denoiser.utils import deserialize_model
from denoiser.dsp import convert_audio

logger = logging.getLogger(__name__)    #日志记录器实例

def add_flags(parser):
    """
    Add the flags for the argument parser that are related to model loading and evaluation"
    """
    parser.add_argument('--device', default="cpu")
    parser.add_argument('--dry', type=float, default=0,
                        help='dry/wet knob coefficient. 0 is only denoised, 1 only input signal.')
    parser.add_argument('--model_path', type=str, default=False,
                        help="Use pre-trained real time H=64 model trained on DNS.")
    parser.add_argument('--model', type=str, default="dns64",  help="Use pre-trained real time H=64 model trained on DNS.")
    parser.add_argument('--num_workers', type=int, default=10)
    parser.add_argument('--streaming', action="store_true",
                        help="true streaming evaluation for Demucs")


parser = argparse.ArgumentParser(
        'denoiser.enhance',
        description="Speech enhancement using Demucs - Generate enhanced files")
add_flags(parser)
parser.add_argument("--out_dir", type=str, default="enhanced_wavs",
                    help="directory putting enhanced wav files")
parser.add_argument("--batch_size", default=1, type=int, help="batch size")
parser.add_argument('-v', '--verbose', action='store_const', const=logging.DEBUG,
                    default=logging.INFO, help="more loggging")

group = parser.add_mutually_exclusive_group()
group.add_argument("--noisy_dir", type=str, default=None,
                   help="directory including noisy wav files")

group.add_argument("--noisy_json", type=str, default=None,
                   help="json file including noisy wav files")

def get_model(args):
    """
    Load local model package or torchhub pre-trained model.
    """
    if args.model_path:
        logger.info("Loading model from %s", args.model_path)
        pkg = torch.load(args.model_path, 'cpu')
        if 'model' in pkg:
            if 'best_state' in pkg:
                pkg['model']['state'] = pkg['best_state']
            model = deserialize_model(pkg['model'])
        else:
            model = deserialize_model(pkg)
    elif args.model == "dns64":
        logger.info("Loading pre-trained real time H=64 model trained on DNS.")
        model = dns64()
    elif args.model == "master64":
        logger.info("Loading pre-trained real time H=64 model trained on DNS and Valentini.")
        model = master64()
    elif args.model == "valentini_nc":
        logger.info("Loading pre-trained H=64 model trained on Valentini.")
        model = valentini_nc()
    else:
        logger.info("Loading pre-trained real time H=48 model trained on DNS.")
        model = dns48()
    logger.debug(model)
    return model

def save_wavs(estimates, noisy_sigs, filenames, out_dir, sr=16_000):
    # Write result
    for estimate, noisy, filename in zip(estimates, noisy_sigs, filenames):
        filename = os.path.join(out_dir, os.path.basename(filename).rsplit(".", 1)[0])
        write(noisy, filename + "_noisy.wav", sr=sr)
        write(estimate, filename + "_enhanced.wav", sr=sr)


def write(wav, filename, sr=16_000):
    # Normalize audio if it prevents clipping
    wav = wav / max(wav.abs().max().item(), 1)
    torchaudio.save(filename, wav.cpu(), sr)

def init_model_denoise():
    args = parser.parse_args()
    model = get_model(args).to(args.device)
    model.eval()
    return model

def denoise(model, wav_path):
    args = parser.parse_args()
    offset = 0
    num_frames = 0
    out, sr = torchaudio.load(str(wav_path))
    target_sr = model.sample_rate or sr
    target_channels = model.chin or out.shape[0]
    out = convert_audio(out, sr, target_sr, target_channels)
    with torch.no_grad():
        estimate = model(out).squeeze(dim=0)
        estimate = (1 - args.dry) * estimate + args.dry * out
    filename = os.path.basename(wav_path).rsplit(".", 1)[0]
    save_file = os.path.join(args.out_dir, filename) + "_enhanced.wav"
    write(estimate, save_file,  sr=model.sample_rate)
    print("denoise finished")


if __name__ == '__main__':
    wav = "record_wav/23_12_14.wav"
    model = init_model_denoise()
    denoise(model, wav)